#!/usr/bin/env python
#coding=UTF-8
'''
Created on 2010-10-23

@author: zarra
'''
from numpy import pi,exp,log2,mat
from numpy.fft import fft

def format_float(f,n):
    fmt="%%.%df"%n
    return fmt%f

class ZFft(object):

    @staticmethod
    def format_complex(c,n=4):
        '''格式化一个复数,精确到n为小数'''
        real = round(c.real,n)
        imag = round(c.imag,n)
        
        if real == imag == 0:
            return '0'
        if imag == 0:
            return format_float(real,n)
        if real == 0:
            return format_float(imag,n)+'j'
        else:
            s='+' if imag >0 else ''
            return "%s%s%sj"%(format_float(real,n),s,format_float(imag,n))
        
    @staticmethod
    def format_complex_list(input):
        '''格式化一复数序列'''
        return ' '.join(ZFft.format_complex(c) for c in input)
    
    @staticmethod
    def bit_reverse(index,n):
        '''码位倒置工具,倒置末n位'''
        result = 0
        for i in range(n):
            b= index &(1<<i)
            
            if b != 0:
                result += (1<<(n-i-1))
        return result
    
    @staticmethod
    def _fft_one_cal(input,size,unit_size,N,Wnx):
        '''进行一次DIF运算
        input:输入序列
        size:序列长度
        unit_size:本次计算段长度
        N:本次为总计算步骤中的第一次
        Wnx:Wnx 序列'''
        
        M = unit_size/2
        mul = 1<<N
        begin = 0
        
        while begin < size:
            for i in range(M):
                x1_index = begin + i
                x2_index = begin + M +i
                
                x1 = input[0,x1_index]
                x2 = input[0,x2_index]
                
                input[0,x1_index] = x1+x2
                input[0,x2_index] = (x1 - x2) * Wnx[i*mul]
            begin += unit_size
                
        
        
    @staticmethod
    def fft(input,size=None):
        '''计算快速傅立叶变换'''
        
        N=0 
        if size ==None:
            size=len(input)
        input=mat([input,input],complex)
        L=int(round(log2(size))) #size的阶数
        Wnx = [ exp(-2j*pi*n/size)for n in range(size/2)]

        unit_size = size
        while (1<<N)  < size:
            ZFft._fft_one_cal(input, size, unit_size, N, Wnx) 
            unit_size = unit_size >> 1
            N+=1
            
        result= list()
        
        
        for i in range(size):
            result.append(input[0,ZFft.bit_reverse(i, L)])
        return result
    
    @staticmethod
    def ifft(input,size=None):
        '''计算快速傅立叶反变换 (调用正变化实现)'''
        if size == None:
            size = len(input)
            
        input = [ c.conjugate() for c in input]
        input = ZFft.fft(input, size)
        
        return [ c.conjugate()/size for c  in input]
    
    @staticmethod
    def fft2d(input_mat,size=None):
        '''计算2维快速傅立叶变换 (调用1维正变化实现)'''
        if size ==None:
            size =input_mat.shape[-1]
        m = mat(input_mat,complex)
        
        for i in range(size):
            t=[m[i,x] for x in range(size)]
            m[i] = ZFft.fft(t,size)          
            
        m = m.T
        
        for i in range(size):
            t=[m[i,x] for x in range(size)]
            m[i] = ZFft.fft(t,size)
            
        return m.T
        
    @staticmethod
    def ifft2d(input_mat,size=None):
        '''计算2维快速傅立叶反变换 (调用二维正变化实现)'''
        if size ==None:
            size =input_mat.shape[-1]
        m=mat(input_mat,complex)
        m=m.conjugate()
        m=ZFft.fft2d(m, size)
        m=m/(size*size)
        return m.conjugate()
        

if __name__ == '__main__':
    import time
    def use_time(f):
        st = time.clock()
        r=f()
        t=time.clock()-st
        return r,t
    
    def dft (seq,N=None):
        if N == None:
            N = len(seq)
        result=[]

        def dft_F(u):
            _sum=0
            for x in range(N):
                _sum+=seq[x]*exp(-2j*pi*u*x/N)
            return _sum 
        
        for u in range(N):
           
            result.append(dft_F(u))
            
        return result

            

    sys_fft    = lambda  x=input:fft(x)
    simple_dft = lambda  x=input:dft(x)
    z_fft      = lambda  x=input:ZFft.fft(x)
    

    input=range(1<<8)
    sys_fft    = lambda  x=input:fft(x)
    simple_dft = lambda  x=input:dft(x)
    z_fft      = lambda  x=input:ZFft.fft(x)
    print use_time(simple_dft)[-1]
    print use_time(sys_fft)[-1]
    print use_time(z_fft)[-1]
    
    input=range(1<<10)
    sys_fft    = lambda  x=input:fft(x)
    simple_dft = lambda  x=input:dft(x)
    z_fft      = lambda  x=input:ZFft.fft(x)
    print use_time(simple_dft)[-1]
    print use_time(sys_fft)[-1]
    print use_time(z_fft)[-1]

    input=range(1<<16)
    sys_fft    = lambda  x=input:fft(x)
    simple_dft = lambda  x=input:dft(x)
    z_fft      = lambda  x=input:ZFft.fft(x)
    print use_time(simple_dft)[-1]
    print use_time(sys_fft)[-1]
    print use_time(z_fft)[-1]
    
    

    
    
    
    